# -*- coding: utf-8 -*-
# 产生批量输入的脚本
# input: (batch_size, 3, seq) 
# output: (batch_size, 6, len(seq)-200)
# 这个可能要定义一个简单的直接读取csv的东西
import config
import os 
import numpy as np
import torch
import copy
from typing import Generator
from matplotlib import pyplot as plt
import sys
sys.path.append("/gpfs/scratch/chgwang/XI/Scripts/MLModel")
import ParallelNet # type: ignore
import matplotlib
matplotlib.use("Agg")

def normalize(data:np.array):
    data = np.array(data)
    max_value = np.max(data)
    min_value = np.min(data)
    if max_value == min_value:
        data = np.ones_like(data)
    else:
        data = (data-min_value) / (max_value - min_value) 
    return(data)

def readFile(path:str):
    # our sample step.
    period = 0.02
    period_points = 200
    sampled_freq = period_points / period
    # this data can got from the header.
    # original sample rate
    with open(path, mode="r") as f:
        while True:
            line = f.readline()
            if "SampleRate" in line:
                break
    line = line.split(",")
    # Sample rate of the data.
    source_freq = float(line[1].strip())
    step = source_freq / sampled_freq
    # convert to int for step only using integer
    step = int(step)
    # get the labels
    path_splited = path.split("/")
    # start with . means is the hidden data
    if path_splited[-1][0] == ".":
        # print(path)
        return
    labels = []
    # "0" as the start name. 
    if path_splited[-1][0] == "0":
        labels.append(["0"])
    else:
        labels.append([path_splited[-1][0]])
        # isdigit is a function for str.m
        if path_splited[-1][1].isdigit():
            labels.append([path_splited[-1][1]])
    # Now the labels is a list of list.
    # read source data.
    # source data shape is (timesteps, channels)
    sour_data = np.loadtxt(path, skiprows=16, delimiter=",", usecols=range(1,4))
    sour_data = sour_data.T
    resampled_data = sour_data[:,::step]
    resample_len = resampled_data.shape[1]
    delimiter = int(resample_len / 2)
    normal_label = ["0"] * delimiter
    for idx, label_list in enumerate(labels):
        label_list = label_list*(resample_len - delimiter)
        temp_list = copy.deepcopy(normal_label)
        temp_list.extend(label_list)
        labels[idx] = temp_list
    # convert labels to int
    labels = np.array(labels).astype(int)
    # norm_data = normalize(resampled_data)
    # norm_data = torch.tensor(norm_data)
    # labels = torch.tensor(labels, dtype=torch.int)
    # return shape is 
    return(resampled_data, labels)

def genPathData(path:str):
    # normalize using in the 
    resampled_data, labels = readFile(path)
    resampled_data_len = resampled_data.shape[1]
    # if not add the "1" the result will be lost 1 sampled
    for i in range(200, resampled_data_len+1):
        sp_data = resampled_data[:,i-200:i]
        sp_data = sp_data.T
        sp_labels = labels[:,i-200:i]
        sp_labels = sp_labels.T
        sp_data = torch.tensor(sp_data)
        sp_data = torch.unsqueeze(sp_data, 0)
        sp_data = torch.unsqueeze(sp_data, 1)
        sp_labels = torch.tensor(sp_labels)
        sp_labels = torch.unsqueeze(sp_labels, 0)
        sp_labels = torch.unsqueeze(sp_labels, 1)
        sp_data = sp_data.float()
        yield(sp_data, sp_labels)

def retrieve_files(path:str) -> Generator:
    path_gen = os.walk(path)
    for root, _, files in path_gen:
        for name in files:
           yield(os.path.join(root,name))

def modeledData(model, data_gen):
    model_out_array = []
    label_out_list = []
    model.eval()
    with torch.no_grad():
        for sp_data, sp_label in data_gen:
            model_out = model(sp_data)
            model_out.cpu().numpy()
            model_out = np.squeeze(model_out)
            assert model_out.shape[0] == 6, \
            "the model_out shapr is %s"%model_out.shape
            model_out_array.append(model_out.cpu().numpy())
            label_out = torch.squeeze(sp_label.cpu()).numpy()
            label_out_list.append(label_out[-1,:])
            # label_out_list.append(torch.squeeze(sp_label.cpu()).numpy())
        model_out_array = np.squeeze(model_out_array)
        label_out_array = np.squeeze(label_out_list)
    return(model_out_array, label_out_array)



def plt_out(model_out_array, label_out_arrry, path_num):
    label_len = label_out_arrry.shape[0]
    delimiter = (label_len + 200) / 2 - 200
    # get the unique labels
    labels = np.unique(label_out_arrry[-1,:]).astype(int)
    if 0 in labels:
        print("This is a noraml experiment")
        return
    else:
        # create subplots
        fig, axes = plt.subplots(label_len, sharex=True)
        fig.suptitle(str(labels))
        for label, ax in zip(labels, axes):
            # the index start with 0
            label = label - 1 # type: int
            ax.axvline(x=delimiter, linestyle="--")
            ax.plot(model_out_array[label,:])
        fig_path = "/gpfs/scratch/chgwang/XI/DataBase/Model_figure/"
        fig.savefig(fig_path+"%s.png"%path_num)

            

if __name__ == "__main__":
    scriptDir = os.path.dirname(os.path.realpath(__file__))
    basepath = os.path.dirname(os.path.dirname(scriptDir))
    dataDir = os.path.join(basepath, "data")
    dataDir = os.path.join(dataDir, "邓茜实验数据")
    file_gen = retrieve_files(dataDir)
    paths = list(file_gen)
    path = paths[0] # type: str
    print(path.split("/")[-1])
    # path data
    data, label = readFile(path)
    # gener a single csv data.
    data_gen = genPathData(path)
    ## TO DO the data input is wrong.
    data, label = next(data_gen)
    # label size of the last two dimension should be 200, 2

    out_dim = 6
    model = ParallelNet.ParallelNet(output_size=out_dim)
    model_path = "/gpfs/scratch/chgwang/XI/DataBase/Model/PN-27-0.734-0.608.pt"
    trained_dict = torch.load(model_path, map_location="cpu")
    model.load_state_dict(trained_dict, strict=True)
    model = model.float()

    modeled_array, labeled_array = modeledData(model,data_gen)
    print(modeled_array.shape)
    print(labeled_array.shape)
    plt_out(modeled_array, labeled_array, 0)
